package main

import (
	_ "RpcNasSync/config/client"
	client2 "RpcNasSync/config/client"
	"RpcNasSync/errno"
	"RpcNasSync/service"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"github.com/juju/ratelimit"
	"github.com/smallnest/rpcx/client"
	"github.com/smallnest/rpcx/share"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
)

type Result struct {
	SP     string `json:"sp"`
	Stdout string `json:"stdout"`
	Stderr string `json:"stderr"`
	Error  string `json:"error"`
}

type Reply struct {
	Stdout string
	Stderr string
	Exit   int
}

type FileOption func(srcFile, fileName, addr string)

var (
	op      = flag.String("op", "", "execute command: shell command or download or upload or uploadDir")
	cmd     = flag.String("cmd", "", "remote server command")
	sp      = flag.String("sp", "", "execute node(L:SP1（SP1/SP2/SP3/SP4/all))")
	file    = flag.String("file", "", "transfer file")
	srcFile = flag.String("src", "", "transfer source file")
	dstFile = flag.String("dst", "", "transfer destination file")
	port    = flag.String("port", "34567", "server port")
	//Conf    = flag.String("conf", "/nasgui/storage/config/nasgui.ini", "nasgui.ini file path")
)

// 执行命令：/nasgui/storage/scripts/nas_exec -op="" -SP=L:SP1（SP1/SP2/SP3/SP4/all)
// 同步文件：/nasgui/storage/scripts/nas_exec -op="upload" file=/root/1.txt SP=L:SP1（SP1/SP2/SP3/SP4/all)  -----相同路径
// /nasgui/storage/scripts/nas_exec -op="download" -src=/root/1.txt -dst=/root/2.txt SP=L:SP1（SP1/SP2/SP3/SP4/all)   -----不同路径
func main() {
	//innerIp, err := innerIp()
	//if err != nil {
	//	fmt.Println("Error:", err)
	//	os.Exit(errno.GetGuiIpErrno)
	//}
	flag.Parse()
	if *sp == "" {
		fmt.Println("Error: Missing required option '-sp'. Please provide a valid node.")
		fmt.Println("Usage: rpc_sync [OPTION]... [FILE]... [NODE]...")
		flag.PrintDefaults()
		return
	}

	if *op == "" && *cmd == "" {
		fmt.Println("Error: Missing required option '-op' or '-cmd'. Please provide a valid operation (-op:shell command/upload/download/uploadDir -cmd:remote server command).")
		fmt.Println("Usage: rpc_sync [OPTION]... [FILE]... [NODE]...")
		flag.PrintDefaults()
		return
	}

	if *cmd != "" {
		processHostFileCmd()
		return
	}

	var fileOp FileOption
	switch *op {
	case "download":
		fileOp = recvFile
	case "upload":
		fileOp = sendFile
	case "uploadDir":
		fileOp = sendDir
	default:
		processHost()
		return
	}

	if *file != "" {
		*srcFile = *file
		*dstFile = *file
	}

	processHostFile(fileOp)

}

func innerIp() ([]string, error) {
	cmd := exec.Command("bash", "-c", "ip a list | grep gui0 | awk '{print $2}' | awk -F '/' '{print $1}'")
	output, _ := cmd.CombinedOutput()
	innerIp := make([]string, 0)
	// 按行分割并过滤空值
	lines := strings.Split(strings.TrimSpace(string(output)), "\n")

	for _, line := range lines {
		if line != "" {
			innerIp = append(innerIp, line)
		}
	}
	if len(innerIp) == 0 {
		return nil, fmt.Errorf("error geting inner0:gui0 ip")
	}
	return innerIp, nil
}

func processHostFile(fileOp FileOption) {
	//fmt.Println(client2.SPNode.Hosts["L:SP1"])
	localhost := client2.SPNode.Hosts["L:SP1"]
	var wg sync.WaitGroup
	if *sp != "all" {
		host := client2.SPNode.Nodes[*sp]
		if localhost == host && *srcFile == *dstFile {
			// 本机节点相同文件,直接返回
			//fmt.Printf("local machine %s and %s are the same object\n", *srcFile, *dstFile)
			return
		}
		if host == "" {
			log.Fatal("Error: Unknown node,Please provide a valid node.")
		}
		fileOp(*srcFile, *dstFile, host)
	} else {
		for _, host := range client2.SPNode.Nodes {
			// 本机节点相同文件,直接返回
			if localhost == host && *srcFile == *dstFile {
				continue
			}
			wg.Add(1)
			go func(k string) {
				defer wg.Done() // 完成时调用 Done()
				fileOp(*srcFile, *dstFile, host)
			}(host)
		}
		wg.Wait()
	}
}

func processHost() {
	resultChan := make(chan Result, 1)
	defer close(resultChan)
	var results []Result
	if *sp != "all" {
		host := client2.SPNode.Nodes[*sp]
		if host == "" {
			log.Fatal("Error: Unknown node,Please provide a valid node.")
		}
		reply, err := execCommand(*op, host)
		if err != nil {
			_, _ = fmt.Fprint(os.Stderr, err)
			var nerr net.Error
			if errors.As(err, &nerr) {
				os.Exit(errno.NetErrno)
			}
			if errors.Is(err, io.ErrUnexpectedEOF) {
				os.Exit(errno.UnexpectedEOF)
			}
		} else {
			fmt.Print(reply.Stdout)
			_, _ = fmt.Fprintf(os.Stderr, reply.Stderr)
			os.Exit(reply.Exit)
		}
	} else {
		for node, host := range client2.SPNode.Nodes {
			go func(node, host string) {
				reply, err := execCommand(*op, host)
				if err != nil {
					resultChan <- Result{SP: node, Stdout: reply.Stdout, Stderr: reply.Stderr, Error: err.Error()}
				} else {
					resultChan <- Result{SP: node, Stdout: reply.Stdout, Stderr: reply.Stderr}
				}
			}(node, host)
		}
		// 收集结果
		for i := 0; i < len(client2.SPNode.Nodes); i++ {
			results = append(results, <-resultChan)
		}
		jsonResult, _ := json.Marshal(results)
		fmt.Println(string(jsonResult))
	}
}

func processHostFileCmd() {
	var wg sync.WaitGroup
	if *sp != "all" {
		host := client2.SPNode.Nodes[*sp]
		if host == "" {
			log.Fatal("Error: Unknown node,Please provide a valid node.")
		}
		err := sendPipeToServe(context.Background(), host, 0, *cmd)
		if err != nil {
			_, _ = fmt.Fprintln(os.Stderr, err)
			os.Exit(errno.SendFileErrno)
		}
	} else {
		for _, host := range client2.SPNode.Nodes {
			wg.Add(1)
			go func(k string) {
				defer wg.Done() // 完成时调用 Done()
				err := sendPipeToServe(context.Background(), host, 0, *cmd)
				if err != nil {
					_, _ = fmt.Fprint(os.Stderr, err)
				}
			}(host)
		}
		wg.Wait()
	}
}

func sendDir(srcDir, dstDir, addr string) {
	d, err := os.ReadDir(srcDir)
	if err != nil {
		_, _ = fmt.Fprintln(os.Stderr, err)
		return
	}
	dirInfo, err := os.Stat(srcDir)
	if err != nil {
		_, _ = fmt.Fprintln(os.Stderr, err)
		return
	}
	reply, err := execCommand(fmt.Sprintf("mkdir -p %s -m %o", dstDir, dirInfo.Mode().Perm()), addr)
	if err != nil {
		_, _ = fmt.Fprintln(os.Stderr, err)
		var nerr net.Error
		if errors.As(err, &nerr) {
			os.Exit(errno.NetErrno)
		}
		return
	}
	if reply.Stderr != "" || reply.Stdout != "" {
		fmt.Print(reply.Stdout)
		_, _ = fmt.Fprintf(os.Stderr, reply.Stderr)
		os.Exit(reply.Exit)
	}
	for _, entry := range d {
		if entry.IsDir() {
			sendDir(filepath.Join(srcDir, entry.Name()), filepath.Join(dstDir, entry.Name()), addr)
		} else {
			sendFile(filepath.Join(srcDir, entry.Name()), filepath.Join(dstDir, entry.Name()), addr)
		}
	}
}

func sendFile(srcFile, fileName, addr string) {
	d, err := client.NewPeer2PeerDiscovery("tcp@"+addr+":"+*port, "")
	if err != nil {
		_, _ = fmt.Fprintf(os.Stderr, "Error: Failed to find server discovery %s: %v\n", addr, err)
		return
	}
	option := client.DefaultOption
	option.RPCPath = share.SendFileServiceName
	xclient := client.NewXClient(share.SendFileServiceName, client.Failfast, client.RandomSelect, d, option)
	defer xclient.Close()
	err = xclient.SendFile(context.Background(), srcFile, 0, map[string]string{"fileName": fileName, "dstIp": addr})
	if err != nil {
		//fmt.Printf("Error: Failed to call %s: %v\n", addr, err)
		_, _ = fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(errno.SendFileErrno)
		return
	}
	//fmt.Printf("Upload to %s file %s success, name %s\n", addr, srcFile, fileName)
}

func recvFile(srcFile, fileName, addr string) {
	d, err := client.NewPeer2PeerDiscovery("tcp@"+addr+":"+*port, "")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: Failed to find server discovery %s: %v\n", addr, err)
		return
	}
	option := client.DefaultOption
	option.RPCPath = share.SendFileServiceName
	xclient := client.NewXClient(share.SendFileServiceName, client.Failfast, client.RandomSelect, d, option)
	defer xclient.Close()
	//dstFile, _ := os.OpenFile(fileName, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0744)
	//defer dstFile.Close()
	//err = xclient.DownloadFile(context.Background(), srcFile, dstFile, map[string]string{"fileName": srcFile})
	remove := true
	err = xclient.DownloadFile2(context.Background(), srcFile, fileName, map[string]string{"fileName": srcFile, "dstIp": addr}, &remove)
	if err != nil {
		if remove == true {
			//删除dstFile
			_ = os.Remove(fileName)
		}
		//fmt.Printf(" %s: %v\n", addr, err)
		_, _ = fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(errno.RecvFileErrno)
	}
	//fmt.Printf("DownLoad from %s file %s success , name: %s\n", addr, srcFile, fileName)
}

func execCommand(cmdString, addr string) (reply Reply, err error) {
	option := client.DefaultOption
	d, err := client.NewPeer2PeerDiscovery("tcp@"+addr+":"+*port, "")
	if err != nil {
		//fmt.Printf("Error:Failed to find server discovery %s: %v\n", addr, err)
		return reply, err
	}
	xclient := client.NewXClient("ServiceShell", client.Failfast, client.RandomSelect, d, option)
	defer xclient.Close()
	args := service.ServiceShell{
		Command: cmdString,
	}
	err = xclient.Call(context.Background(), "ShellCommand", args, &reply)
	if err != nil {
		//fmt.Printf("Error: Failed to call %s: %v\n", addr, err)
		//fmt.Printf("%v\n", err)
		return reply, err
	}
	return reply, nil
}

func sendPipeToServe(ctx context.Context, addr string, rateInBytesPerSecond int64, cmd string) error {
	d, err := client.NewPeer2PeerDiscovery("tcp@"+addr+":"+*port, "")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: Failed to find server discovery %s: %v\n", addr, err)
		return err
	}
	xclient := client.NewXClient(share.SendFileServiceName, client.Failfast, client.RandomSelect, d, client.DefaultOption)
	// 默认文件名可以从 meta 或其他地方获取
	args := share.FileTransferArgs{
		FileName: cmd, // 默认文件名标识
		FileSize: -1,  // 流式传输无法预先知道大小，设为 -1
	}
	// 从标准输入中读取数据
	file := os.Stdin

	defer file.Close()

	ctx = context.Background()
	reply := &share.FileTransferReply{}
	err = xclient.Call(ctx, "TransferFile", args, reply)
	if err != nil {
		fmt.Fprintln(os.Stderr, "Error: Failed to Call serviceMethod TransferFile")
		return err
	}
	reply.Addr = addr + reply.Addr
	conn, err := net.Dial("tcp", reply.Addr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: Failed to Dial %s\n", reply.Addr)
		return err
	}

	defer conn.Close()

	_, err = conn.Write(reply.Token)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: Failed to Write Token %s\n", reply.Token)
		return err
	}

	var tb *ratelimit.Bucket

	if rateInBytesPerSecond > 0 {
		tb = ratelimit.NewBucketWithRate(float64(rateInBytesPerSecond), rateInBytesPerSecond)
	}

	sendBuffer := make([]byte, client.FileTransferBufferSize)
	readBuffer := make([]byte, client.FileTransferBufferSize)
loop:
	for {
		select {
		case <-ctx.Done():
			err = ctx.Err()
			break loop
		default:
			if tb != nil {
				tb.Wait(client.FileTransferBufferSize)
			}
			n, err := file.Read(sendBuffer)
			if err != nil {
				if err == io.EOF {
					return nil
				} else {
					return err
				}
			}
			if n == 0 {
				break loop
			}
			_, err = conn.Write(sendBuffer[:n])
			if err != nil {
				if err == io.EOF {
					return nil
				} else {
					conn.Read(readBuffer)
					//fmt.Println(string(readBuffer[:n]))
					return fmt.Errorf(string(readBuffer))
				}
			}
		}
	}

	return nil
}
